import time
import numpy as np
import cv2


def cone_plane_cross(Q, plane):
    """
    求解圆锥与平面的交点
    cone: 圆锥
    plane: 平面
    """
    # 提取平面方程的参数
    a, b, c, d = plane
    if len(Q.shape) == 2:
        Q = Q[None, :, :]

    # 计算交线的齐次矩阵 Q'
    A_prime = Q[:, 0, 0] * c**2 - 2 * a * c * Q[:, 0, 2] + (a**2) * Q[:, 2, 2]
    E_prime = Q[:, 1, 1] * c**2 - 2 * b * c * Q[:, 1, 2] + (b**2) * Q[:, 2, 2]
    B_prime = (
        Q[:, 0, 1] * c**2
        - a * c * Q[:, 1, 2]
        - b * c * Q[:, 0, 2]
        + (a * b) * Q[:, 2, 2]
    )
    D_prime = (
        Q[:, 0, 3] * c**2
        - a * c * Q[:, 2, 3]
        - d * c * Q[:, 0, 2]
        + (a * d) * Q[:, 2, 2]
    )
    G_prime = (
        Q[:, 1, 3] * c**2
        - b * c * Q[:, 2, 3]
        - d * c * Q[:, 1, 2]
        + (b * d) * Q[:, 2, 2]
    )
    J_prime = Q[:, 3, 3] * c**2 - 2 * d * c * Q[:, 2, 3] + (d**2) * Q[:, 2, 2]

    return np.array(
        [
            [A_prime, B_prime, D_prime],
            [B_prime, E_prime, G_prime],
            [D_prime, G_prime, J_prime],
        ]
    ).T.squeeze()


def ransac_ellipse_fit(x, y, t=6, k=100, th=10):
    sample = np.random.choice(range(len(x)), t, replace=False)
    best_num = 0
    best_choice = None
    x_homo = np.block([[x], [y], [np.ones_like(x)]])
    for _ in range(k):
        C = ellipse_fit(x[sample], y[sample])
        err = np.abs(x_homo * (C @ x_homo)).sum(0) / x_homo[2] ** 2
        inliers = err < th
        if inliers.sum() > best_num:
            best_num = inliers.sum()
            best_choice = inliers
            if best_num > t:
                sample = np.arange(len(x))[inliers]
            else:
                sample = np.random.choice(range(len(x)), t, replace=False)
        else:
            sample = np.random.choice(range(len(x)), t, replace=False)
    if best_choice is None:
        return None
    else:
        return ellipse_fit(x[best_choice], y[best_choice])


def ellipse_fit(x, y):
    x = x[:, np.newaxis]
    y = y[:, np.newaxis]
    # A = np.hstack((2 * x, 2 * y, np.ones_like(x)))
    # X = np.linalg.lstsq(A, x * x + y * y, rcond=None)[0].squeeze()
    # A = np.array(
    #     [
    #         [1, 0, -X[0]],
    #         [0, 1, -X[1]],
    #         [-X[0], -X[1], -X[2]],
    #     ]
    # )
    A = np.hstack([x * x, 2 * x * y, y * y, 2 * x, 2 * y, np.ones_like(x)])
    # X = np.linalg.lstsq(A, -x * x, rcond=None)[0].squeeze()
    _, _, V = np.linalg.svd(A)
    X = V[-1, :]
    A = np.array(
        [
            [X[0], X[1], X[3]],
            [X[1], X[2], X[4]],
            [X[3], X[4], X[5]],
        ]
    )
    return A


def from_ellipse_args(x, y, a, b, theta):
    x = (
        x
        + a * np.cos(np.linspace(0, 2 * np.pi, 200)) * np.cos(theta)
        - b * np.sin(np.linspace(0, 2 * np.pi, 200)) * np.sin(theta)
    )
    y = (
        y
        + a * np.cos(np.linspace(0, 2 * np.pi, 200)) * np.sin(theta)
        + b * np.sin(np.linspace(0, 2 * np.pi, 200)) * np.cos(theta)
    )
    return np.vstack((x, y))


def center_ellipse(C):
    x_cnt = (C[0, 1] * C[1, 2] - C[1, 1] * C[0, 2]) / (
        C[0, 0] * C[1, 1] - C[0, 1] * C[0, 1]
    )
    y_cnt = (C[0, 1] * C[0, 2] - C[0, 0] * C[1, 2]) / (
        C[0, 0] * C[1, 1] - C[0, 1] * C[0, 1]
    )
    return x_cnt, y_cnt
    # return -np.linalg.inv(C[:2, :2]) @ C[:2, 2]


def ellipse_args(C):
    x = center_ellipse(C)
    # g = np.linalg.det(C) / np.linalg.det(C[:2, :2])
    # eigvals, eigvecs = np.linalg.eig(C[:2, :2])
    # a = np.sqrt(-g / eigvals[0])
    # b = np.sqrt(-g / eigvals[1])
    # theta = np.arctan2(
    #     2 * C[0, 1],
    #     C[0, 0] - C[1, 1] - np.sqrt((C[0, 0] - C[1, 1]) ** 2 + 4 * C[0, 1] ** 2),
    # )

    # return *x, a, b, theta
    a, b, c, d, e, f = C[0, 0], C[0, 1], C[1, 1], C[0, 2], C[1, 2], C[2, 2]
    val1 = (
        2
        * (a * e * e - 2 * b * d * e + c * d * d - (a * c - b * b) * f)
        / (a * c - b * b)
        / (a + c + np.sqrt((a - c) ** 2 + 4 * b * b))
    )
    val2 = (
        2
        * (a * e * e - 2 * b * d * e + c * d * d - (a * c - b * b) * f)
        / (a * c - b * b)
        / (a + c - np.sqrt((a - c) ** 2 + 4 * b * b))
    )
    if val1 < 0 or val2 < 0:
        return None, None, None, None, None
    r1 = np.sqrt(val1)
    r2 = np.sqrt(val2)
    theta = np.arctan2(2 * b, a - c - np.sqrt((a - c) ** 2 + 4 * b * b))
    return *x, r1, r2, theta


def pose_calculate(K, X, x):
    distCoeffs = np.zeros(4)
    # 计算旋转向量和平移向量
    success, rvec, tvec = cv2.solvePnP(X, x, K, distCoeffs)
    if not success:
        Warning("Failed to solve PnP")
        return None, None
    # 将旋转向量转换为旋转矩阵
    R, _ = cv2.Rodrigues(rvec)
    return R, tvec


def ransac_pnp(K, points_3d, observations, th=8):
    best_num = 0
    best_inliers = np.array([False] * len(observations))
    sample = np.random.choice(len(observations), 4, replace=False)
    for _ in range(100):
        R, tvec = pose_calculate(K, points_3d[sample], observations[sample])
        if R is not None:
            err = np.linalg.norm(
                cv2.projectPoints(points_3d, R, tvec, K, np.zeros(4))[0].squeeze()
                - observations,
                axis=1,
            )
            inliers = err < th
            if inliers.sum() > best_num:
                best_num = inliers.sum()
                best_inliers = inliers
            if inliers.sum() > 0.8 * len(observations):
                break
            # 另一种方式，使用绝对点个数，实际上超过20个点后计算所得精度已经相当高了。
            # if inliers.sum() > 20:
            #     break
            if inliers.sum() > 4:
                sample = np.random.choice(np.where(inliers)[0], 4, replace=False)
            else:
                sample = np.random.choice(len(observations), 4, replace=False)
        else:
            sample = np.random.choice(len(observations), 4, replace=False)
    if best_inliers.sum() < 4:
        return None
    else:
        return best_inliers


class PoseEstimator:
    def __init__(self, K, *args, **kwargs):
        """
        必须的输入参数：
        K: 相机内参，3x3矩阵
        """
        self.K = K

    def __call__(self, database, measurements, *args, **kwargs):
        """
        必须的输入参数：
        database: 数据库存储的数据结构，如三维点、空间平面的椭圆等
        measurements: 观测到的数据，如二维点、像平面上的椭圆等
        kwargs: 其他可选参数
        """
        start = time.time()
        result = self.forward(database, measurements, *args, **kwargs)
        end = time.time()
        return end - start, *result

    def forward(self, K, database, measurements, *args, **kwargs):
        raise NotImplementedError

    def __name__(self):
        raise NotImplementedError

    def __repr__(self):
        return (
            f"<class {self.__name__()} > ("
            + ", ".join(f"{k}={v}" for k, v in self.__dict__.items() if k != "K")
            + ")"
        )
